import re
import pymysql
import datetime
from utils.llm import ask_deepseek as llm_func

class AutoSQL:
    def __init__(self, host, port, user, password, database, llm_func, max_len=2048, max_try=1):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.llm_func = llm_func
        self.max_len = max_len
        self.max_try = max_try
        # Connect to the database
        self.conn = pymysql.connect(host=self.host, port=self.port,
                                    user=self.user, password=self.password,
                                    database=self.database)
        # Create a cursor
        self.cursor = self.conn.cursor()
        # Print all the tables in the database
        self.cursor.execute("show tables")
        tables = self.cursor.fetchall()
        print("Tables in the database:")
        for table in tables:
            print(table[0])
    
    def get_all_tables(self):
        table_names = []
        all_columns = []
        self.cursor.execute("show tables")
        tables = self.cursor.fetchall()
        for table in tables:
            table_names.append(table[0])
            # get the columns of the table
            table_name = table[0]
            sql = "show full columns from %s" % table_name
            self.cursor.execute(sql)
            results = self.cursor.fetchall()
            for result in results:
                field = result[0]
                if field not in all_columns:
                    all_columns.append(field)
        return table_names, all_columns
    
    def get_table_info(self,tablename):
        # 获取表的字段信息以及备注信息的描述字符串
        info = ""
        sql = "show full columns from %s" % tablename
        self.cursor.execute(sql)
        results = self.cursor.fetchall()
        index = 1
        all_columns = []
        for result in results:
            field = result[0]
            comment = result[8]
            info += f"列名{index}: {field}, 说明:{comment}\n"
            index += 1
            all_columns.append(field)
        return info, all_columns
    

    def analysis_instruction(self, instruction, tablename, tableinfo, history=None):
        # 用大模型分析指令，给出查询步骤和逻辑，比如先查询到什么信息，然后根据新的信息去查询什么信息
        if (not tableinfo) or (not tablename):
            query_str = ""
            table_index = 1
            for tablename in self.get_all_tables()[0]:
                tableinfo,all_columns = self.get_table_info(tablename)
                query_str += f"\n<tableinfo{table_index}>\n{tableinfo}\n</tableinfo{table_index}>\n <tablename{table_index}>\n{tablename}\n</tablename{table_index}>"
                table_index += 1
            query_str += f"\n<instruction>\n{instruction}\n</instruction>"
        else:
            query_str = f"\n<tableinfo>\n{tableinfo}\n</tableinfo>\n<tablename>\n{tablename}\n</tablename>\n<instruction>\n{instruction}\n</instruction>"
        query_str += "\nPlease generate the SQL statement according to the table information and the instruction."
        query_str += "\n<return example>SELECT * FROM table WHERE condition</return>"
        query_str += "\n<note>Never use = , use like %value% instead.</note>"
        query_str += "\n<note>Use more () to avoid the priority problem</note>"
        query_str += "\n<note>岗位名称，人名，部门名称等具体value，都是中文的</note>"
        query_str += "<note>如果使用日期，日期不要使用=，因为日期是精确到秒的。SQL中不要有任何备注。<note>"
        query_str += "\n<note>使用的表格越少越好，如非必要，不要使用join联查</note>"
        query_str += "<task>请你根据instruction以及其他信息，整理出查询的步骤和逻辑</task>\n实例如下："
        query_str += "<输入实例>总结赵胜的部门员工上周的工作呢绒</输入实例>"
        query_str += """<返回示例>
```json
{
    'tasks': ['获取姓名等于赵胜的部门','查询上一步骤中部门的所有员工','查询上一步骤中员工的最近一周日志'],
    'logics': ['根据员工姓名查询员工所在部门','根据部门名称查询部门员工','根据员工姓名查询员工最近一周的日志']
}
```
</返回示例>"""
        text = self.llm_func(query_str)
        # 提取```json ```中包裹的部分，并转化成字典格式
        import re
        result = re.search(r"```json(.*?)```", text, re.DOTALL)
        # 获取match的text，然后用eval转为字典
        try:
            result = eval(result.group(1))
            print(f"LLM return: {result}")
            tasks = result["tasks"]
            logics = result["logics"]
            return tasks, logics
        except Exception as e:
            print(f"Error in analysis_instruction: {e}")
            return [], []
    
    # 根据analysis_instruction返回的tasks和logic，逐步生成SQL语句完善查询结果，返回完整的查询结果
    def get_info_from_tasks(self, tasks, logics, tablename, tableinfo, history=None):
        knowledges = []
        executed_sqls = []
        for i in range(len(tasks)):
            task = tasks[i]
            logic = logics[i]
            
            print(f"Task: {task}, Logic: {logic}")
            if i == 0:
                # 第一个任务，直接生成SQL语句
                sql = self.intelligent_generate_sql(tablename, tableinfo, f"{task} {logic}", history)
                correct,results = self.run_sql(sql)
                if not correct:
                    break
                r = self.format_result(results)
                knowledges.append(r)
                executed_sqls.append(sql)
                
            else:
                # 之后的任务，根据上一步的结果，生成SQL语句
                sql = self.intelligent_generate_sql(tablename, tableinfo, f"{task} {logic}, 上一步的sql为：{sql} 上一步的输出信息为：{r}", history)
                correct,results = self.run_sql(sql)
                if not correct:
                    break
                r = self.format_result(results)
                knowledges.append(r)
                executed_sqls.append(sql)
        result = ""
        for i in range(len(tasks)):
            try:
                result += f"Task: {tasks[i]}, Logic: {logics[i]}\n"
                result += f"Knowledge: {knowledges[i]}\n"
            except Exception as e:
                return result
        return result

    def intelligent_generate_sql(self, tablename, tableinfo, instruction, history=None):
        if (not tableinfo) or (not tablename):
            query_str = ""
            table_index = 1
            for tablename in self.get_all_tables()[0]:
                tableinfo,all_columns = self.get_table_info(tablename)
                query_str += f"\n<tableinfo{table_index}>\n{tableinfo}\n</tableinfo{table_index}>\n <tablename{table_index}>\n{tablename}\n</tablename{table_index}>"
                table_index += 1
            query_str += f"\n<instruction>\n{instruction}\n</instruction>"
        else:
            query_str = f"\n<tableinfo>\n{tableinfo}\n</tableinfo>\n<tablename>\n{tablename}\n</tablename>\n<instruction>\n{instruction}\n</instruction>"
        query_str += "\nPlease generate the SQL statement according to the table information and the instruction."
        query_str += "\n<return example>SELECT * FROM table WHERE condition</return>"
        query_str += "\n<note>Never use = , use like %value% instead.</note>"
        query_str += "\n<note>Use more () to avoid the priority problem</note>"
        query_str += "\n<note>岗位名称，人名，部门名称等具体value，都是中文的</note>"
        query_str += "如果使用日期，日期永远不要使用=，只能用>=或者<=。SQL中不要有任何备注。"
        query_str += "\n<note>使用的表格越少越好，如非必要，不要使用join联查</note>"

        # query_str += "<note> if you think there is no need to use the table, just return ''</note>"
        text = self.llm_func(query_str, history=history)
        print(f"LLM return: {text}")
        # Get the generated SQL statement
        if text.startswith('"') and text.endswith('"'):
            text = text[1:-1]
        # 如果存在```sql```标记，则提取其中的SQL语句
        if "```sql" in text:
            text = text.split("```sql")[1].split("```")[0]
        text = text.replace('"',"'")
        # if WHERE xxx LIKE 'xxx' in the SQL statement, replace 'xxx' with %xxx%
        
        if "LIKE" in text:
            # 高级替换，一步到位
            import re
            text = re.sub(r"LIKE\s+'([^']*)'", r"LIKE '%\1%'", text)
        # 把 xx = xxx 替换为 xx LIKE %xxx%
        print(text)
        return text
    
    def run_sql(self, sql, all_columns=[]):
        try:
            sql = sql.replace("\n", " ")
            self.cursor.execute(sql)
            results = self.cursor.fetchall()
            if len(results) == 0:
                # Try repalce the collumn to another collumn until get the result
                for column in all_columns:
                    if column in sql:
                        for column2 in all_columns:
                            if column2 != column and "date" not in column2 and "time" not in column2:
                                sql2 = sql.replace(column, column2)
                                print(f"Try to replace {column} to {column2}")
                                self.cursor.execute(sql2)
                                results = self.cursor.fetchall()
                                if len(results) > 0:
                                    return True,results
                return False,"No result"
                
            return True,results
        except Exception as e:
            return False,str(e)
        
    def format_result(self, results):
        print(f"Results: {results}")
        result_str = ""
        for result in results:
            already_added_info = []
            for item in result:
                # 如果是日期或者时间类型，转换为字符串
                # 打印item的类型
                print(type(item))
                if isinstance(item, datetime.datetime) or\
                    isinstance(item, datetime.date):
                    item = item.strftime("%Y-%m-%d")
                if item in already_added_info:
                    continue
                else:
                    already_added_info.append(item)
                if not item:
                    continue
                result_str += str(item).replace("\r\n","").replace("\n","")[:100] + "\t"
            result_str += "\n"
            if len(result_str) > self.max_len:
                break
        return result_str
    
    def close(self):
        self.cursor.close()
        self.conn.close()

    def auto_run(self, instruction, tablename=None, add_date=True, history=None):
        if add_date:
            instruction += " " + f"现在是{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
        if tablename:
            tableinfo,all_columns = self.get_table_info(tablename)
        else:
            print(f"No table name, try to get all tables")
            tableinfo = None
            tablename = None
            all_columns = self.get_all_tables()[1]
        try_times = 0
        while try_times < self.max_try:
            sql = self.intelligent_generate_sql(tablename, tableinfo, instruction, history)
            correct,results = self.run_sql(sql, all_columns)
            if not correct:
                try_times += 1
                instruction += f"\n<note>\n{sql} is error, {results}\n</note>\n"
                print(f"Try times: {try_times}")
            else:
                break
        if not correct:
            return ""
        return self.format_result(results)

# Initialize the AutoSQL object
autosql = AutoSQL(
    host="106.14.113.100",
    port=12316,
    user="root",
    password="longyuan",
    database="dingding",
    # llm func return input
    llm_func=llm_func
)
# Run the SQL statement

def ask_sql(instruct, tablename=None, history=None, max_history=5):
    
    if tablename:
        tableinfo,all_columns =autosql.get_table_info(tablename)
    else:
        tableinfo = None
        tablename = None
    tasks,logics = autosql.analysis_instruction(instruct, tablename, tableinfo, history)
    other_info_mation = autosql.get_info_from_tasks(tasks, logics, tablename, tableinfo, history)
    print(other_info_mation)
    

    if history:
        if len(history) > max_history:
            history = history[-max_history:]
    results = autosql.auto_run(instruct, tablename=tablename, history=history)
    return results+f"\n迭代查询结果：{other_info_mation}"

if __name__ == "__main__":
    ask_sql("java岗位都有谁？")